#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
股价与情感指数可视化分析工具
用于分析股票价格与舆论情感之间的关系
"""

import os
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.ticker import MaxNLocator
import numpy as np
from datetime import datetime
import seaborn as sns
# from scipy import stats
import matplotlib as mpl
import platform
from matplotlib.font_manager import FontProperties
plt.rcParams['font.sans-serif'] = ['SimHei']
# Matplotlib中设置字体-黑体，解决Matplotlib中文乱码问题
plt.rcParams['axes.unicode_minus'] = False
# 解决Matplotlib坐标轴负号'-'显示为方块的问题
sns.set(font='SimHei')

font_path = './font/SimHei.ttf'
font_prop = FontProperties(fname=font_path)
plt.rcParams['font.family'] = font_prop.get_name()
mpl.font_manager.fontManager.addfont(font_path)

class StockSentimentVisualizer:
    """股价与情感指数可视化分析类"""
    
    def __init__(self, stock_data_path, sentiment_data_path):
        """
        初始化可视化器
        
        参数:
            stock_data_path: 股价数据CSV文件路径
            sentiment_data_path: 情感指数CSV文件路径
        """
        self.stock_data_path = stock_data_path
        self.sentiment_data_path = sentiment_data_path
        self.stock_data = None
        self.sentiment_data = None
        self.merged_data = None
        
        # 加载数据
        self._load_data()
        
    def _load_data(self):
        """加载股价和情感指数数据"""
        # 检查文件是否存在
        if not os.path.exists(self.stock_data_path):
            raise FileNotFoundError(f"股价数据文件不存在: {self.stock_data_path}")
        
        if not os.path.exists(self.sentiment_data_path):
            raise FileNotFoundError(f"情感指数数据文件不存在: {self.sentiment_data_path}")
        
        # 加载股价数据
        self.stock_data = pd.read_csv(self.stock_data_path)
        print(f"成功加载股价数据，共 {len(self.stock_data)} 条记录")
        
        # 加载情感指数数据
        self.sentiment_data = pd.read_csv(self.sentiment_data_path)
        print(f"成功加载情感指数数据，共 {len(self.sentiment_data)} 条记录")
        
        # 确保日期列格式正确
        self._format_date_columns()
        
        # 合并数据
        self._merge_data()
    
    def _format_date_columns(self):
        """处理日期列格式"""
        # 对股价数据处理日期列
        if "日期" in self.stock_data.columns:
            self.stock_data.rename(columns={"日期": "date"}, inplace=True)
        
        # 确保日期列格式统一
        try:
            self.stock_data['date'] = pd.to_datetime(self.stock_data['date'])
            self.sentiment_data['date'] = pd.to_datetime(self.sentiment_data['date'])
        except Exception as e:
            print(f"日期格式转换错误: {str(e)}")
            # 如果转换失败，尝试另一种方式
            try:
                self.stock_data['date'] = pd.to_datetime(self.stock_data['date'], format='%Y-%m-%d')
                self.sentiment_data['date'] = pd.to_datetime(self.sentiment_data['date'], format='%Y-%m-%d')
            except Exception as e:
                print(f"第二次尝试日期格式转换错误: {str(e)}")
    
    def _merge_data(self):
        """合并股价和情感指数数据"""
        # 确保两个数据集都有日期列
        if 'date' not in self.stock_data.columns or 'date' not in self.sentiment_data.columns:
            raise ValueError("股价数据或情感指数数据缺少date列")
        
        # 处理股价数据中的列名
        if "收盘" in self.stock_data.columns:
            self.stock_data.rename(columns={"收盘": "close_price"}, inplace=True)
        elif "Close" in self.stock_data.columns:
            self.stock_data.rename(columns={"Close": "close_price"}, inplace=True)
        
        # 使用日期列合并数据
        self.merged_data = pd.merge(
            self.stock_data[['date', 'close_price']], 
            self.sentiment_data[['date', 'sentiment_index']], 
            on='date', 
            how='inner'
        )
        
        # 将日期设置为索引
        self.merged_data.set_index('date', inplace=True)
        
        # 按日期排序
        self.merged_data.sort_index(inplace=True)
        self.merged_data = self.merged_data[-125:]
        print(f"数据合并完成，共 {len(self.merged_data)} 条记录")
        print(self.merged_data.head())
    
    def plot_trend(self, title="股价与情感指数走势图", figsize=(14, 8), save_path=None, dpi=100):
        """
        绘制股价和情感指数的走势图
        
        参数:
            title: 图表标题
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) == 0:
            print("没有可用的数据进行绘图")
            return
        
        # 创建图表
        fig, ax1 = plt.subplots(figsize=figsize)
        
        # 设置主题样式
        # plt.style.use('seaborn-v0_8-whitegrid')
        
        # 绘制股价曲线
        color = 'tab:red'
        ax1.set_xlabel('日期')
        ax1.set_ylabel('股价', color=color)
        ax1.plot(self.merged_data.index, self.merged_data['close_price'], color=color, linewidth=2, marker='o', label='股价')
        ax1.tick_params(axis='y', labelcolor=color)
        
        # 设置x轴日期格式
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        
        # 创建第二个y轴
        ax2 = ax1.twinx()
        color = 'tab:blue'
        ax2.set_ylabel('情感指数', color=color)
        ax2.plot(self.merged_data.index, self.merged_data['sentiment_index'], color=color, linewidth=2, marker='x', label='情感指数')
        ax2.tick_params(axis='y', labelcolor=color)
        
        # 添加网格线
        ax1.grid(True, linestyle='--', alpha=0.6)
        
        # 添加图例
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        # 设置标题
        plt.title(title, fontsize=15)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"图表已保存至: {save_path}")
        
        # 显示图表
        plt.show()
    
    def plot_correlation(self, figsize=(14, 10), save_path=None, dpi=100):
        """
        绘制股价和情感指数的相关性分析图
        
        参数:
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) < 5:
            print("数据不足，无法进行相关性分析")
            return
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        
        # 设置主题样式
        # plt.style.use('seaborn-v0_8-whitegrid')
        
        # 1. 计算相关系数
        correlation = self.merged_data['close_price'].corr(self.merged_data['sentiment_index'])
        corr_text = f"相关系数: {correlation:.4f}"
        
        # 2. 散点图和回归线
        sns.regplot(
            x='sentiment_index', 
            y='close_price', 
            data=self.merged_data, 
            ax=axes[0, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        axes[0, 0].set_title(f'股价与情感指数散点图 ({corr_text})')
        axes[0, 0].set_xlabel('情感指数')
        axes[0, 0].set_ylabel('股价')
        
        # 3. 热力图
        sns.heatmap(
            self.merged_data[['close_price', 'sentiment_index']].corr(),
            annot=True,
            cmap='coolwarm',
            ax=axes[0, 1],
            vmin=-1, vmax=1
        )
        axes[0, 1].set_title('相关性热力图')
        
        # 4. 股价变化与情感指数的关系
        # 计算每日股价变化百分比
        self.merged_data['price_change_pct'] = self.merged_data['close_price'].pct_change() * 100
        
        # 去除NaN值
        data_change = self.merged_data.dropna()
        
        # 散点图和回归线
        change_corr = data_change['price_change_pct'].corr(data_change['sentiment_index'])
        change_corr_text = f"相关系数: {change_corr:.4f}"
        
        sns.regplot(
            x='sentiment_index', 
            y='price_change_pct', 
            data=data_change, 
            ax=axes[1, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        axes[1, 0].set_title(f'股价变化百分比与情感指数 ({change_corr_text})')
        axes[1, 0].set_xlabel('情感指数')
        axes[1, 0].set_ylabel('股价变化百分比 (%)')
        
        # 5. 情感指数分布直方图
        sns.histplot(self.merged_data['sentiment_index'], kde=True, ax=axes[1, 1])
        axes[1, 1].set_title('情感指数分布')
        axes[1, 1].set_xlabel('情感指数')
        axes[1, 1].set_ylabel('频率')
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"相关性分析图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
    
    def lag_analysis(self, max_lag=10, figsize=(14, 6), save_path=None, dpi=100):
        """
        滞后相关性分析
        
        参数:
            max_lag: 最大滞后天数
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) <= max_lag:
            print("数据不足，无法进行滞后分析")
            return
        
        # 计算不同滞后期的相关系数
        correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                # 当前相关性
                corr = self.merged_data['close_price'].corr(self.merged_data['sentiment_index'])
            else:
                # 滞后相关性 (股价滞后于情感)
                corr = self.merged_data['close_price'].corr(self.merged_data['sentiment_index'].shift(-i))
            correlations.append(corr)
        
        # 创建图表
        plt.figure(figsize=figsize)

        # plt.style.use('seaborn-v0_8-whitegrid')
        
        plt.bar(range(max_lag + 1), correlations, color='skyblue', alpha=0.7)
        plt.plot(range(max_lag + 1), correlations, 'ro-', linewidth=2)
        
        plt.xlabel('滞后天数')
        plt.ylabel('相关系数')
        plt.title('股价与情感指数的滞后相关性分析')
        plt.xticks(range(max_lag + 1))
        plt.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        
        # 添加数值标签
        for i, v in enumerate(correlations):
            plt.text(i, v + 0.02 * (1 if v >= 0 else -1), f'{v:.3f}', ha='center')
        
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"滞后相关性分析图已保存至: {save_path}")
        
        plt.show()
        
        # 找出最大相关系数对应的滞后期
        max_corr_idx = np.argmax(np.abs(correlations))
        max_corr = correlations[max_corr_idx]
        print(f"最强相关性在滞后 {max_corr_idx} 天时出现，相关系数为 {max_corr:.4f}")
        
        return correlations

    def add_smoothed_sentiment(self):
        """添加平滑处理后的情感指数"""
        windows = [3, 7, 14]  # 不同时间窗口
        for window in windows:
            self.merged_data[f'sentiment_ma{window}'] = self.merged_data['sentiment_index'].rolling(window=window).mean()
        
        # 也可以尝试指数移动平均
        self.merged_data['sentiment_ema7'] = self.merged_data['sentiment_index'].ewm(span=7).mean()

    def time_scale_analysis(self, save_path=None):
        """多尺度时间分析"""
        # 按周/月聚合数据
        weekly_data = self.merged_data.resample('W').mean()
        monthly_data = self.merged_data.resample('ME').mean()  # 使用'ME'代替'M'
        
        # 绘制不同时间尺度下的关系图
        fig, axes = plt.subplots(3, 1, figsize=(14, 18))
        
        # 日尺度
        self._plot_dual_axis(self.merged_data, axes[0], '日频数据')
        # 周尺度
        self._plot_dual_axis(weekly_data, axes[1], '周频数据')
        # 月尺度
        self._plot_dual_axis(monthly_data, axes[2], '月频数据')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path)
        plt.show()

    def cross_correlation_analysis(self, max_lag=30):
        """交叉相关分析，找出最佳领先/滞后关系"""
        from statsmodels.tsa.stattools import ccf
        
        # 标准化数据
        price_norm = (self.merged_data['close_price'] - self.merged_data['close_price'].mean()) / self.merged_data['close_price'].std()
        sentiment_norm = (self.merged_data['sentiment_index'] - self.merged_data['sentiment_index'].mean()) / self.merged_data['sentiment_index'].std()
        
        # 计算交叉相关
        ccf_result = ccf(price_norm.dropna(), sentiment_norm.dropna(), adjusted=False)
        
        # 确保我们只使用有效的CCF结果范围
        valid_lags = min(max_lag, len(ccf_result))
        half_lags = valid_lags // 2
        
        # 获取实际有效的CCF结果
        ccf_values = ccf_result[:valid_lags]
        
        # 创建对应的横坐标
        x_range = list(range(-half_lags, valid_lags - half_lags))
        
        # 确保x和y长度相同
        if len(x_range) > len(ccf_values):
            x_range = x_range[:len(ccf_values)]
        elif len(ccf_values) > len(x_range):
            ccf_values = ccf_values[:len(x_range)]
        
        # 找出最大相关及其滞后期
        max_corr_idx = np.argmax(np.abs(ccf_values))
        lag = x_range[max_corr_idx]
        
        plt.figure(figsize=(12, 6))
        plt.bar(x_range, ccf_values, alpha=0.7)
        plt.axvline(x=lag, color='r', linestyle='--')
        plt.title(f'股价与情感指数交叉相关分析 (最大相关滞后期: {lag}天)')
        plt.xlabel('滞后天数 (负值表示情感指数领先股价)')
        plt.ylabel('相关系数')
        plt.grid(True, alpha=0.3)
        plt.show()

    def segment_analysis(self, n_segments=3):
        """将时间序列分成多个区段进行分析"""
        # 将数据分成n_segments个时间段
        dates = self.merged_data.index
        segment_size = len(dates) // n_segments
        
        plt.figure(figsize=(15, 10))
        
        for i in range(n_segments):
            start_idx = i * segment_size
            end_idx = (i+1) * segment_size if i < n_segments-1 else len(dates)
            
            segment_data = self.merged_data.iloc[start_idx:end_idx]
            corr = segment_data['close_price'].corr(segment_data['sentiment_index'])
            
            plt.subplot(n_segments, 1, i+1)
            plt.plot(segment_data.index, segment_data['close_price'], 'r-', label='股价')
            plt.title(f'时间段 {i+1}: {segment_data.index[0].date()} 至 {segment_data.index[-1].date()}, 相关系数={corr:.4f}')
            
            plt2 = plt.twinx()
            plt2.plot(segment_data.index, segment_data['sentiment_index'], 'b-', label='情感指数')
            
            # 合并图例
            lines1, labels1 = plt.gca().get_legend_handles_labels()
            lines2, labels2 = plt2.get_legend_handles_labels()
            plt2.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
        
        plt.tight_layout()
        plt.show()

    def granger_causality_test(self, max_lag=10):
        """格兰杰因果检验，分析情感指数是否对股价有预测作用"""
        from statsmodels.tsa.stattools import grangercausalitytests
        
        # 准备数据
        data = pd.DataFrame({
            'price': self.merged_data['close_price'],
            'sentiment': self.merged_data['sentiment_index']
        })
        data = data.dropna()
        
        # 情感指数是否对股价有格兰杰因果关系
        print("检验情感指数是否'引起'股价变动:")
        grangercausalitytests(data[['price', 'sentiment']], max_lag, verbose=True)
        
        # 股价是否对情感指数有格兰杰因果关系
        print("\n检验股价是否'引起'情感指数变动:")
        grangercausalitytests(data[['sentiment', 'price']], max_lag, verbose=True)

    def volatility_correlation(self, window=7):
        """分析股价波动率与情感指数波动性的关系"""
        # 计算股价的波动率(用简单的滚动标准差)
        self.merged_data['price_volatility'] = self.merged_data['close_price'].rolling(window).std()
        
        # 计算情感指数的波动性
        self.merged_data['sentiment_volatility'] = self.merged_data['sentiment_index'].rolling(window).std()
        
        # 去除NaN
        vol_data = self.merged_data.dropna()
        
        # 计算相关系数
        corr = vol_data['price_volatility'].corr(vol_data['sentiment_volatility'])
        
        # 绘制散点图
        plt.figure(figsize=(10, 8))
        plt.scatter(vol_data['sentiment_volatility'], vol_data['price_volatility'], alpha=0.6)
        plt.title(f'股价波动率与情感指数波动性相关分析 (相关系数: {corr:.4f})')
        plt.xlabel('情感指数波动性')
        plt.ylabel('股价波动率')
        
        # 添加回归线
        z = np.polyfit(vol_data['sentiment_volatility'], vol_data['price_volatility'], 1)
        p = np.poly1d(z)
        plt.plot(vol_data['sentiment_volatility'], p(vol_data['sentiment_volatility']), "r--")
        
        plt.grid(True, alpha=0.3)
        plt.show()

    def extreme_sentiment_event_analysis(self, threshold=30):
        """极端情感事件对股价的影响分析"""
        # 找出情感指数极高和极低的日期
        high_sentiment = self.merged_data[self.merged_data['sentiment_index'] > threshold].index
        low_sentiment = self.merged_data[self.merged_data['sentiment_index'] < -threshold].index
        
        # 分析这些日期后n天的股价表现
        windows = [1, 3, 5, 10]
        
        results = pd.DataFrame(index=['极高情感后', '极低情感后'], columns=[f'{w}天后平均涨跌幅(%)' for w in windows])
        
        for window in windows:
            high_returns = []
            low_returns = []
            
            for date in high_sentiment:
                try:
                    future_date = self.merged_data.index[self.merged_data.index.get_loc(date) + window]
                    pct_change = (self.merged_data.loc[future_date, 'close_price'] / 
                                 self.merged_data.loc[date, 'close_price'] - 1) * 100
                    high_returns.append(pct_change)
                except:
                    pass
                
            for date in low_sentiment:
                try:
                    future_date = self.merged_data.index[self.merged_data.index.get_loc(date) + window]
                    pct_change = (self.merged_data.loc[future_date, 'close_price'] / 
                                 self.merged_data.loc[date, 'close_price'] - 1) * 100
                    low_returns.append(pct_change)
                except:
                    pass
            
            results.loc['极高情感后', f'{window}天后平均涨跌幅(%)'] = np.mean(high_returns) if high_returns else np.nan
            results.loc['极低情感后', f'{window}天后平均涨跌幅(%)'] = np.mean(low_returns) if low_returns else np.nan
        
        return results

    def _plot_dual_axis(self, data, ax, title):
        """
        在指定的轴上绘制股价和情感指数的双轴图
        
        参数:
            data: 包含股价和情感指数的DataFrame
            ax: 要绘制的matplotlib轴对象
            title: 图表标题
        """
        # 绘制股价曲线
        color1 = 'tab:red'
        ax.set_xlabel('日期')
        ax.set_ylabel('股价', color=color1)
        ax.plot(data.index, data['close_price'], color=color1, linewidth=2, marker='o', label='股价')
        ax.tick_params(axis='y', labelcolor=color1)
        
        # 设置x轴日期格式
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
        
        # 创建第二个y轴
        ax2 = ax.twinx()
        color2 = 'tab:blue'
        ax2.set_ylabel('情感指数', color=color2)
        ax2.plot(data.index, data['sentiment_index'], color=color2, linewidth=2, marker='x', label='情感指数')
        ax2.tick_params(axis='y', labelcolor=color2)
        
        # 添加网格线
        ax.grid(True, linestyle='--', alpha=0.6)
        
        # 添加图例
        lines1, labels1 = ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        # 设置标题
        ax.set_title(title)
        
        # 计算相关系数
        corr = data['close_price'].corr(data['sentiment_index'])
        ax.annotate(f"相关系数: {corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction', 
                   fontsize=10, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))

def main():
    """主函数"""
    # 文件路径
    stock_data_path = "../../data/stock_data/hk01810_stock_data(3).csv"
    sentiment_data_path = "../../data/sentiment_index2.csv"
    
    # 创建可视化器
    visualizer = StockSentimentVisualizer(stock_data_path, sentiment_data_path)
    
    # 绘制走势图
    visualizer.plot_trend(title="小米集团(01810)股价与情感指数走势图",
                          save_path="result2/stock_sentiment_trend.png")
    
    # 绘制相关性分析图
    visualizer.plot_correlation(save_path="result2/stock_sentiment_correlation.png")
    
    # 进行滞后分析
    visualizer.lag_analysis(max_lag=5, save_path="result2/stock_sentiment_lag_analysis.png")
    visualizer.time_scale_analysis()
    visualizer.cross_correlation_analysis()
    # visualizer.granger_causality_test()
    print(visualizer.extreme_sentiment_event_analysis().head())
if __name__ == "__main__":
    main() 